from __future__ import print_function
import tensorflow as tf
import keras
from keras import regularizers
from keras import backend as K
from keras.layers import Input,Conv1D,Dense,Flatten,Reshape,BatchNormalization
from tdcnn_layers import mask_Conv1D,fix_Conv1D
from keras.models import Model
import numpy as np
import pickle as pkl
import pandas as pd
import os


pwd = os.getcwd()
[flow,relate_cross,crossName,target_cross,pred_cross,similar_cross,roadNet,submitList,data_tr]=pkl.load(open(pwd+"/data.pkl",'rb'))


def normalize(data, mean_min, std_max, norm_method="gussian"):
    if norm_method == "gussian":
        return (data - mean_min) / std_max
    elif norm_method == "linear":
        return (data - mean_min) / (std_max - mean_min)
    
def de_normalize(data, mean_min, std_max, norm_method="gussian"):
    if norm_method == "gussian":
        return data * std_max + mean_min
    elif norm_method == "linear":
        return data * (std_max - mean_min) + mean_min
    
def root_mean_squared_error(y_true, y_pred):
        return K.sqrt(K.mean(K.square(y_pred - y_true), axis=-1))
    
def change_data_form(cross):
    #history,   space,      time,      similirity,  label
    timeNo,flow_day,flow_minute,flow_y = data_tr[cross]
    flow_minute = flow_minute.values
    flow_day = flow_day.values
    timeNo = timeNo.values
    flow_y = flow_y.values
    
    total_num = len(timeNo)
    tr_num = int(total_num * 0.8)
    seed = 3
    np.random.seed(seed)
    train_indices = np.random.choice(total_num, tr_num, replace=False)# 设置随机数生成从0-150中随机挑选120个随机数
    test_indices = np.array(list(set(range(total_num)) - set(train_indices)))
    tr_mnt = flow_minute[train_indices]
    tr_day = flow_day[train_indices]
    tr_time = timeNo[train_indices]
    tr_y = flow_y[train_indices]
    val_mnt = flow_minute[test_indices]
    val_day = flow_day[test_indices]
    val_time = timeNo[test_indices]
    val_y = flow_y[test_indices]
        
    mean_min = np.mean(tr_y)# traffic flow
    std_max = np.std(tr_y)
    data = [tr_y,val_y,tr_day,val_day,tr_mnt,val_mnt,tr_time,val_time,mean_min,std_max]
    
    for i in range(6):
        data[i] = normalize(data[i], mean_min, std_max)
        data[i] = data[i].astype(np.float32)
    
    for i in range(2,6):
        data[i] =  data[i].reshape(list(data[i].shape)+[1])
        
    return data

def get_batch_size(epoch):
    if epoch < 2:
        batch_size = 16
    elif epoch < 4:
        batch_size = 64
    elif epoch < 6:
        batch_size = 128
    else:
        batch_size = 256
    return batch_size

def circle_reg(weight_matrix):#相邻向量差别不能过大，第一个和最后一个差别也不能过大
    print( weight_matrix.shape )
    [rows, cols] = weight_matrix.shape
    rows = int(rows)
    loss = 0
    for ii in range(rows):
        loss += K.sqrt(K.square(weight_matrix[ii%rows]-weight_matrix[ (ii+1)%rows ]))
    return 0.01 * loss

def build_model():
    day_shape = (7,1)
    minute_shape = (6,1) # minute, point, channel(flow-speed)
    
    K.set_learning_phase(True) #set learning phase
    # fix random seed for reproducibility,初始化模型时会有一些随机数,不固定,不同次初始化,即使经过加载参数,结果略有不同
    seed = 7
    np.random.seed(seed)
    # inputXY, X->layer-th, Y:1->space info, 2->period info
    input_day = Input(shape=day_shape, dtype='float32', name='input_day')
    input_mnt = Input(shape=minute_shape, dtype='float32', name='input_mnt')
    input_time = Input(shape=(1,), dtype='int32', name='time')
    onehot = tf.one_hot(input_time, 23)
    dense22 = Dense(6,activation='relu',name='dense22',kernel_regularizer=circle_reg)(onehot)
    dense22 = BatchNormalization()(dense22)
    dense22 = Reshape([6,1])(dense22)
    dense22 = fix_Conv1D(1, kernel_size=(1),use_bias=False,activity_regularizer=regularizers.l1(0.01),
                       kernel_initializer=tf.constant_initializer(0.3))(dense22)
    dense22 = Reshape([6,1,1])(dense22)
        
    conv2_day = Conv1D(16,kernel_size=(3),activation='relu')(input_day)
    conv2_mnt = mask_Conv1D(16,kernel_size=(3),padding='same',activation='relu',mode='knn')(input_mnt,inputs2=dense22,masks=5)
    
    conv3_day = Conv1D(32,kernel_size=(3,),activation='relu',input_shape=(8,5,3,16))(conv2_day)
    conv3_mnt = Conv1D(32,kernel_size=(3),activation='relu',input_shape=(8,8,16))(conv2_mnt)
    
    flat4_day = Flatten()(conv3_day)
    flat4_mnt = Flatten()(conv3_mnt)    
    con5 = keras.layers.concatenate([flat4_day, flat4_mnt])
    
    dense6 = Dense(512, activation='relu', name='dense6',kernel_regularizer=regularizers.l2(0.01))(con5)
    output = Dense(6, name='output',activation='relu')(dense6)
    
    model = Model(inputs=[input_day, input_mnt, input_time], outputs=output)
    model.compile(optimizer='adam',loss=root_mean_squared_error,metrics=['mse'])#rmsprop
    return model
    
    
def train_model(cross):
    model = build_model()
    tr_y,val_y,tr_day,val_day,tr_mnt,val_mnt,tr_time,val_time,mean_min,std_max=change_data_form(cross)    
    min_rmse_real = 999999
    for epoch in range(8):
        batch_size = get_batch_size(epoch)
        epochs = 1
        model.fit({'input_day':tr_day,'input_mnt':tr_mnt,'time':tr_time},tr_y,
                  epochs=epochs,batch_size=batch_size,verbose=1,
                  validation_data=({'input_day':val_day,'input_mnt':val_mnt,'time':val_time},val_y))
        K.set_learning_phase(False) #set learning phase
        score = model.evaluate({'input_day':val_day,'input_mnt':val_mnt,'time':val_time},
                               val_y,verbose=0,batch_size=batch_size)
        
        pred = model.predict({'input_day':val_day,'input_mnt':val_mnt,'time':val_time},batch_size=2)
        #pred=pred.reshape(-1)
        
        pred_real = de_normalize(pred, mean_min, std_max)
        val_y_real = de_normalize(val_y, mean_min, std_max)
        
        mae_real = np.mean(np.abs(pred_real - val_y_real))
        mse_real = np.mean(pow(pred_real - val_y_real, 2))
        mape_real = np.mean(np.abs(pred_real - val_y_real)/val_y_real)
        rmse_real = np.sqrt(np.mean(pow(pred_real - val_y_real, 2)))
        mse_norm = np.mean(pow(pred - val_y, 2))
        print('mape_real:',mape_real,'mse_real:',mse_real,'rmse_real:',rmse_real,'mse_norm:',mse_norm,"epoch:",epoch)
        print('val loss:',score[0],' val accuracy:',score[1],'mae_real:',mae_real)
        print('=== rmse_real ===:',rmse_real)
        if rmse_real < min_rmse_real:
            min_rmse_real = rmse_real
            model.save_weights(pwd+"/model/"+str(cross)+".h5")
            print('----------------------- min_rmse_real ---------------------------:',min_rmse_real)
            # 保存预测结果
            tmp_df = pd.DataFrame(pred_real,val_y_real)
            tmp_df.to_csv(pwd+"/val_predict/"+str(cross)+".csv",sep=':',header=None)



# ------------------------------- train model for pred_cross ----------------------------------
import sys
sys.exit(0)


for cross in pred_cross:
    train_model(cross)
    sys.exit(0)

print('-------- train done. --------')

# --------------------------------- predict pred_cross -----------------------------------------

#count = 0
tmp_day = list(range(7))
tmp_minute = list(range(6))
testTime = {7,9,11,14,16,18} # No of time_range
model = build_model()
model.compile(optimizer='adam',loss=root_mean_squared_error,metrics=['mse'])
for cross in pred_cross:    
    model.load_weights(pwd+"/model/"+str(cross)+".h5")
    count = 0
    timeNo = pd.DataFrame(np.zeros([1,1]))
    flow_day = pd.DataFrame(np.zeros([1,7]))
    flow_minute = pd.DataFrame(np.zeros([1,6]))
    # flow_y = pd.DataFrame(np.zeros([1,6])) --> pd.DataFrame(np.zeros([6,1]))
    for k in range(5478,len(flow)):
        time_ = flow.index[k]
        time_tmp = int( int(time_.hour)*100 + int(time_.minute) )
        if time_.day >= 20 and time_.day <= 23 and time_.minute == 30 and time_.hour in testTime:
            print(cross,': ',time_)
            timeNo.loc[count] = time_.hour*2 + (1 if time_.minute>=30 else 0) - 15
            for ii in range(1,8):
                #print(k-288*ii)
                tmp_day[ii-1] = flow.iloc[k-288*ii][cross]
            flow_day.loc[count] = tmp_day
            for ii in range(1,7):
                tmp_minute[ii-1] = flow.iloc[k-ii][cross]
            flow_minute.loc[count] = tmp_minute
            #count = count + 1
            input_day = flow_day.values.reshape(list(flow_day.values.shape)+[1])
            input_mnt = flow_minute.values.reshape(list(flow_minute.values.shape)+[1])
            flow_y = model.predict({'input_day':input_day,'input_mnt':input_mnt,'time':timeNo.values})
            pred_rng = pd.date_range(start=time_, periods=6, freq = '5min')
            #flow_y = pd.DataFrame(np.transpose(flow_y.values), index=pred_rng, columns=[cross])
            flow.loc[pred_rng,cross] = flow_y.reshape(-1)
            #for time_stamp in pred_rng:
                #print(time_stamp)
            #    flow.loc[time_stamp, cross] = flow_y.loc[time_stamp,'result']
            
pkl.dump(flow,open(pwd+"/result/acnn_result.pkl",'wb'))


            

    












